#include <planner/jps_planner/jps_graph_search.h>
#include <cmath>
#include <iostream>

using namespace sz;

JPSGraphSearch::JPSGraphSearch(const uint8_t *cMap, int xDim, int yDim, float eps, bool verbose) :
    cMap_(cMap), xDim_(xDim), yDim_(yDim), eps_(eps), verbose_(verbose) {
  hm_.resize(xDim_ * yDim_);
  seen_.resize(xDim_ * yDim_, false);

  for (int x = -1; x <= 1; x++) {
    for (int y = -1; y <= 1; y++) {
      if (x == 0 && y == 0) continue;
      ns_.emplace_back(std::vector<int>{x, y});
    }
  }

}

inline int JPSGraphSearch::coordToId(short x, short y) const {
  return x + y * xDim_;
}

inline bool JPSGraphSearch::isFree(int x, int y) const {
  return x >= 0 && x < xDim_ && y >= 0 && y < yDim_ && cMap_[coordToId(x, y)] == val_free_;
}

inline bool JPSGraphSearch::isOccupied(int x, int y) const {
  return x >= 0 && x < xDim_ && y >= 0 && y < yDim_ && cMap_[coordToId(x, y)] != val_free_;
}

inline float JPSGraphSearch::getHeur(int x, int y) const {
  return eps_ * std::sqrt((x - xGoal_) * (x - xGoal_) + (y - yGoal_) * (y - yGoal_));
}

bool JPSGraphSearch::plan(short xStart, short yStart, short xGoal, short yGoal, int maxExpand) {
  priority_queue_.clear();
  path_.clear();
  hm_.resize(xDim_ * yDim_);
  seen_.resize(xDim_ * yDim_, false);


  // Set goal
  int goal_id = coordToId(xGoal, yGoal);
  xGoal_ = xGoal;
  yGoal_ = yGoal;

  // Set start node
  int start_id = coordToId(xStart, yStart);
  JPSStatePtr currNode_ptr = std::make_shared<JPSState>(JPSState(start_id, xStart, yStart, 0, 0));
  currNode_ptr->g = 0;
  currNode_ptr->h = getHeur(xStart, yStart);

  return plan(currNode_ptr, maxExpand, start_id, goal_id);
}

bool JPSGraphSearch::plan(JPSStatePtr &currNode_ptr, int maxExpand, int start_id, int goal_id) {
  // Insert start node
  currNode_ptr->handle_type = priority_queue_.push(currNode_ptr);
  currNode_ptr->opened = true;
  hm_[currNode_ptr->id] = currNode_ptr;
  seen_[currNode_ptr->id] = true;

  int expand_iteration = 0;
  while (true) {
    expand_iteration++;
    // get element with smallest cost
    currNode_ptr = priority_queue_.top();
    priority_queue_.pop();
    currNode_ptr->closed = true; // Add to closed list

    if (currNode_ptr->id == goal_id) {
      if (verbose_)
        printf("Goal Reached!!!!!!\n\n");
      break;
    }

    //printf("expand: %d, %d\n", currNode_ptr->x, currNode_ptr->y);
    std::vector<int> succ_ids;
    std::vector<float> succ_costs;
    // Get successors
    getJpsSucc(currNode_ptr, succ_ids, succ_costs);

//    if (verbose_)
//      printf("size of succs: %zu\n", succ_ids.size());
    // Process successors
    for (int s = 0; s < (int) succ_ids.size(); s++) {
      //see if we can improve the value of succstate
      JPSStatePtr &child_ptr = hm_[succ_ids[s]];
      float tentative_gval = currNode_ptr->g + succ_costs[s];

      if (tentative_gval < child_ptr->g) {
        child_ptr->parentId = currNode_ptr->id;  // Assign new parent
        child_ptr->g = tentative_gval;    // Update gval


        // if currently in OPEN, update
        if (child_ptr->opened && !child_ptr->closed) {
          priority_queue_.increase(child_ptr->handle_type);       // update heap
          child_ptr->dx = child_ptr->x - currNode_ptr->x;
          child_ptr->dy = child_ptr->y - currNode_ptr->y;
          if (child_ptr->dx != 0)
            child_ptr->dx /= std::abs(child_ptr->dx);
          if (child_ptr->dy != 0)
            child_ptr->dy /= std::abs(child_ptr->dy);
        } else if (child_ptr->opened && child_ptr->closed) { // if currently in CLOSED
          printf("ASTAR ERROR!\n");
        } else { // new node, add to heap
          //printf("add to open set: %d, %d\n", child_ptr->x, child_ptr->y);
          child_ptr->handle_type = priority_queue_.push(child_ptr);
          child_ptr->opened = true;
        }
      } //
    } // Process successors


    if (maxExpand > 0 && expand_iteration >= maxExpand) {
      if (verbose_)
        printf("MaxExpandStep [%d] Reached!!!!!!\n\n", maxExpand);
      return false;
    }

    if (priority_queue_.empty()) {
      if (verbose_)
        printf("Priority queue is empty!!!!!!\n\n");
      return false;
    }
  }

  if (verbose_) {
    printf("goal g: %f, h: %f!\n", currNode_ptr->g, currNode_ptr->h);
    printf("Expand [%d] nodes!\n", expand_iteration);
  }

  recoverPath(currNode_ptr, start_id);

  return true;
}

void JPSGraphSearch::recoverPath(JPSStatePtr node, int start_id) {
  VectorPoint2i path;
  path.emplace_back(node->x, node->y);
  while (node && node->id != start_id) {
    node = hm_[node->parentId];
    path.emplace_back(node->x, node->y);
  }
  std::reverse(std::begin(path), std::end(path));
  path_.swap(path);
}

void JPSGraphSearch::getJpsSucc(const JPSStatePtr &curr, std::vector<int> &succ_ids, std::vector<float> &succ_costs) {

  const int norm1 = std::abs(curr->dx) + std::abs(curr->dy);
  int8_t num_neib = jn2d_.nsz[norm1][0];
  int8_t num_fneib = jn2d_.nsz[norm1][1];
  int id = (curr->dx + 1) + 3 * (curr->dy + 1);

  for (int dev = 0; dev < num_neib + num_fneib; ++dev) {
    short new_x, new_y;
    int8_t dx, dy;
    if (dev < num_neib) {
      dx = jn2d_.ns[id][0][dev];
      dy = jn2d_.ns[id][1][dev];
      if (!jump(curr->x, curr->y, dx, dy, new_x, new_y)) continue;
    } else {
      int nx = curr->x + jn2d_.f1[id][0][dev - num_neib];
      int ny = curr->y + jn2d_.f1[id][1][dev - num_neib];
      if (isOccupied(nx, ny)) {
        dx = jn2d_.f2[id][0][dev - num_neib];
        dy = jn2d_.f2[id][1][dev - num_neib];
        if (!jump(curr->x, curr->y, dx, dy, new_x, new_y)) continue;
      } else
        continue;
    }

    int new_id = coordToId(new_x, new_y);
    if (!seen_[new_id]) {
      seen_[new_id] = true;
      hm_[new_id] = std::make_shared<JPSState>(new_id, new_x, new_y, dx, dy);
      hm_[new_id]->h = getHeur(new_x, new_y);
    }

    succ_ids.push_back(new_id);
    succ_costs.push_back(std::sqrt((new_x - curr->x) * (new_x - curr->x) +
        (new_y - curr->y) * (new_y - curr->y)));
  }

}

bool JPSGraphSearch::jump(short x, short y, short dx, short dy, short &new_x, short &new_y) {
  new_x = x + dx;
  new_y = y + dy;
  if (!isFree(new_x, new_y))
    return false;

  if (new_x == xGoal_ && new_y == yGoal_)
    return true;

  if (hasForced(new_x, new_y, dx, dy))
    return true;

  const int id = (dx + 1) + 3 * (dy + 1);
  const short norm1 = std::abs(dx) + std::abs(dy);
  int8_t num_neib = jn2d_.nsz[norm1][0];
  for (int8_t k = 0; k < num_neib - 1; ++k) {
    short new_new_x, new_new_y;
    if (jump(new_x, new_y, jn2d_.ns[id][0][k], jn2d_.ns[id][1][k], new_new_x, new_new_y))
      return true;
  }

  return jump(new_x, new_y, dx, dy, new_x, new_y);
}

inline bool JPSGraphSearch::hasForced(short x, short y, short dx, short dy) {
  const int id = (dx + 1) + 3 * (dy + 1);
  for (int8_t fn = 0; fn < 2; ++fn) {
    if (isOccupied(x + jn2d_.f1[id][0][fn], y + jn2d_.f1[id][1][fn]))
      return true;
  }
  return false;
}

const VectorPoint2i JPSGraphSearch::getPath() const {
  return path_;
}

std::vector<JPSStatePtr> JPSGraphSearch::getOpenSet() const {
  std::vector<JPSStatePtr> ss;
  for (const auto &it: hm_) {
    if (it && it->opened && !it->closed)
      ss.emplace_back(it);
  }
  return ss;
}

std::vector<JPSStatePtr> JPSGraphSearch::getCloseSet() const {
  std::vector<JPSStatePtr> ss;
  for (const auto &it: hm_) {
    if (it && it->closed)
      ss.push_back(it);
  }
  return ss;
}

std::vector<JPSStatePtr> JPSGraphSearch::getAllSet() const {
  std::vector<JPSStatePtr> ss;
  for (const auto &it: hm_) {
    if (it)
      ss.push_back(it);
  }
  return ss;
}

constexpr int8_t JPS2DNeib::nsz[3][2];

JPS2DNeib::JPS2DNeib() {
  int id = 0;
  for (int8_t dy = -1; dy <= 1; ++dy) {
    for (int8_t dx = -1; dx <= 1; ++dx) {
      short norm1 = std::abs(dx) + std::abs(dy);
      for (int8_t dev = 0; dev < nsz[norm1][0]; ++dev)
        Neib(dx, dy, norm1, dev, ns[id][0][dev], ns[id][1][dev]);
      for (int8_t dev = 0; dev < nsz[norm1][1]; ++dev) {
        FNeib(dx, dy, norm1, dev,
              f1[id][0][dev], f1[id][1][dev],
              f2[id][0][dev], f2[id][1][dev]);
      }
      id++;
    }
  }
}

void JPS2DNeib::print() {
  for (int8_t dx = -1; dx <= 1; dx++) {
    for (int8_t dy = -1; dy <= 1; dy++) {
      int id = (dx + 1) + 3 * (dy + 1);
      printf("[dx: %d, dy: %d]-->id: %d:\n", dx, dy, id);
      for (unsigned int i = 0; i < sizeof(f1[id][0]) / sizeof(f1[id][0][0]); i++)
        printf("                f1: [%d, %d]\n", f1[id][0][i], f1[id][1][i]);
    }
  }
}

void JPS2DNeib::Neib(int8_t dx, int8_t dy, int8_t norm1, int8_t dev, int8_t &tx, int8_t &ty) {
  switch (norm1) {
    case 0:
      switch (dev) {
        case 0: tx = 1;
          ty = 0;
          return;
        case 1: tx = -1;
          ty = 0;
          return;
        case 2: tx = 0;
          ty = 1;
          return;
        case 3: tx = 1;
          ty = 1;
          return;
        case 4: tx = -1;
          ty = 1;
          return;
        case 5: tx = 0;
          ty = -1;
          return;
        case 6: tx = 1;
          ty = -1;
          return;
        case 7: tx = -1;
          ty = -1;
          return;
        default:return;
      }
    case 1:tx = dx;
      ty = dy;
      return;
    case 2:
      switch (dev) {
        case 0: tx = dx;
          ty = 0;
          return;
        case 1: tx = 0;
          ty = dy;
          return;
        case 2: tx = dx;
          ty = dy;
          return;
        default:return;
      }
    default:return;
  }
}

void JPS2DNeib::FNeib(int8_t dx, int8_t dy, int8_t norm1, int8_t dev, int8_t &fx, int8_t &fy, int8_t &nx, int8_t &ny) {
  switch (norm1) {
    case 1:
      switch (dev) {
        case 0: fx = 0;
          fy = 1;
          break;
        case 1: fx = 0;
          fy = -1;
          break;
        default:break;
      }

      // switch order if different direction
      if (dx == 0)
        fx = fy, fy = 0;

      nx = dx + fx;
      ny = dy + fy;
      return;
    case 2:
      switch (dev) {
        case 0:fx = -dx;
          fy = 0;
          nx = -dx;
          ny = dy;
          return;
        case 1:fx = 0;
          fy = -dy;
          nx = dx;
          ny = -dy;
          return;
        default:return;
      }
    default:return;
  }
}